import os
import unittest

import torch
from torch import hub
from torch.backends import cudnn
from torchsummary import summary

from config import Config
from training.models import fasige, fasige_2, fasige_2_only_gen, fasige_2_only_dcm

# CONFIG = Config()
# hub.set_dir(CONFIG['TORCH_HOME'])
# os.environ["CUDA_VISIBLE_DEVICES"] = CONFIG['CUDA_VISIBLE_DEVICES']
#
# torch.backends.cudnn.benchmark = True

if torch.cuda.is_available() == False:
    # CPU:
    CONFIG = Config()
    hub.set_dir(CONFIG['WINDOWS_TORCH_HOME'])
else:
    # GPU:
    CONFIG = Config()
    hub.set_dir(CONFIG['TORCH_HOME'])
    os.environ["CUDA_VISIBLE_DEVICES"] = CONFIG['CUDA_VISIBLE_DEVICES']
    torch.backends.cudnn.benchmark = True


class SifdNetTestCase(unittest.TestCase):

    def test_fasige(self):
        # self.assertTrue(torch.cuda.is_available())
        # model = sifdnet_1(pretrained=True)
        # model = fasige(pretrained=True)
        model = fasige_2(pretrained=True)
        # model = fasige_2_only_gen(pretrained=True)
        # model = fasige_2_only_dcm(pretrained=True)

        if not torch.cuda.is_available():
            model = model.cpu()
        else:
            model = model.cuda()
        input_size = model.default_cfg['input_size']
        summary(model, input_size=input_size)


if __name__ == '__main__':
    unittest.main()
